"""
极化码核心算法实现
包括编码、信道构造和Fast-SCL译码
"""

import numpy as np
from typing import List, Tuple, Optional

class PolarCode:
    """极化码实现类"""
    
    def __init__(self, N: int, K: int, L: int = 8):
        """
        初始化极化码
        
        Args:
            N: 码长 (必须是2的幂次)
            K: 信息位长度
            L: SCL译码列表大小
        """
        self.N = N
        self.K = K
        self.L = L
        self.n = int(np.log2(N))
        
        # 生成极化信道可靠性序列
        self.frozen_bits, self.info_bits = self._generate_frozen_bits()
        
    def _generate_frozen_bits(self) -> Tuple[List[int], List[int]]:
        """生成冻结位和信息位的索引"""
        # 使用高斯近似计算信道可靠性
        z = np.zeros(self.N)
        
        # 初始化
        z[0] = 0.5
        
        # 递归计算
        for level in range(self.n):
            step = 2 ** level
            for i in range(0, self.N, 2 * step):
                for j in range(step):
                    if i + j < self.N:
                        z[i + j] = 2 * z[i + j] - (z[i + j] ** 2)
                        if i + j + step < self.N:
                            z[i + j + step] = z[i + j] ** 2
        
        # 排序得到最可靠的信道
        reliability = np.argsort(z)
        
        # 选择K个最可靠的作为信息位
        info_bits = reliability[:self.K].tolist()
        frozen_bits = reliability[self.K:].tolist()
        
        return frozen_bits, info_bits
    
    def encode(self, message: np.ndarray) -> np.ndarray:
        """
        极化码编码
        
        Args:
            message: 信息位向量 (长度K)
            
        Returns:
            编码后的码字 (长度N)
        """
        if len(message) != self.K:
            raise ValueError(f"消息长度应为{self.K}，但得到{len(message)}")
            
        # 初始化输入向量
        u = np.zeros(self.N)
        
        # 填充信息位
        for i, bit_idx in enumerate(self.info_bits):
            u[bit_idx] = message[i]
        
        # 极化变换
        return self._polar_transform(u)
    
    def _polar_transform(self, u: np.ndarray) -> np.ndarray:
        """极化变换"""
        x = u.copy()
        
        # 递归进行极化变换
        for stage in range(self.n):
            step = 2 ** (self.n - stage - 1)
            for i in range(0, self.N, 2 * step):
                for j in range(step):
                    if i + j + step < self.N:
                        temp = x[i + j]
                        x[i + j] = (x[i + j] + x[i + j + step]) % 2
                        x[i + j + step] = temp
        
        return x
    
    def decode(self, received: np.ndarray, snr_db: float) -> np.ndarray:
        """
        Fast-SCL译码
        
        Args:
            received: 接收到的信号
            snr_db: 信噪比(dB)
            
        Returns:
            译码结果
        """
        snr = 10 ** (snr_db / 10)
        sigma = np.sqrt(1 / (2 * snr))
        
        # 计算LLR
        llr = 2 * received / (sigma ** 2)
        
        # Fast-SCL译码
        return self._fast_scl_decode(llr)
    
    def _fast_scl_decode(self, llr: np.ndarray) -> np.ndarray:
        """Fast-SCL译码算法实现"""
        # 初始化路径
        paths = [{'bits': [], 'metric': 0.0, 'active': True}]
        
        # 逐位译码
        for bit_idx in range(self.N):
            new_paths = []
            
            # 对每个活跃路径进行扩展
            for path in paths:
                if not path['active']:
                    continue
                    
                # 计算当前位的LLR
                llr_val = self._calculate_llr(llr, path['bits'], bit_idx)
                
                # 如果是冻结位
                if bit_idx in self.frozen_bits:
                    new_bit = 0
                    new_metric = path['metric'] + self._bit_metric(llr_val, new_bit)
                    new_paths.append({
                        'bits': path['bits'] + [new_bit],
                        'metric': new_metric,
                        'active': True
                    })
                else:
                    # 信息位，分叉
                    for bit_val in [0, 1]:
                        new_metric = path['metric'] + self._bit_metric(llr_val, bit_val)
                        new_paths.append({
                            'bits': path['bits'] + [bit_val],
                            'metric': new_metric,
                            'active': True
                        })
            
            # 路径修剪
            paths = self._prune_paths(new_paths)
            
            # Fast-SCL优化：提前终止
            if len(paths) == 1 and bit_idx < self.N - 1:
                # 检查是否可以提前终止
                if self._can_early_terminate(paths[0], bit_idx):
                    # 完成剩余位的译码
                    remaining_bits = self._complete_decoding(paths[0], bit_idx + 1)
                    paths[0]['bits'].extend(remaining_bits)
                    break
        
        # 选择最佳路径
        best_path = min(paths, key=lambda x: x['metric'])
        
        # 提取信息位
        decoded_bits = []
        for bit_idx in self.info_bits:
            if bit_idx < len(best_path['bits']):
                decoded_bits.append(best_path['bits'][bit_idx])
        
        return np.array(decoded_bits)
    
    def _calculate_llr(self, llr: np.ndarray, decoded_bits: List[int], bit_idx: int) -> float:
        """计算当前位的LLR值"""
        # 简化的LLR计算
        return llr[bit_idx]
    
    def _bit_metric(self, llr: float, bit: int) -> float:
        """计算比特度量"""
        return abs(llr) if (llr < 0 and bit == 1) or (llr >= 0 and bit == 0) else 0
    
    def _prune_paths(self, paths: List[dict]) -> List[dict]:
        """路径修剪"""
        # 按度量排序
        paths.sort(key=lambda x: x['metric'])
        
        # 保留前L个路径
        return paths[:self.L]
    
    def _can_early_terminate(self, path: dict, bit_idx: int) -> bool:
        """检查是否可以提前终止"""
        # 简化的提前终止条件
        return True
    
    def _complete_decoding(self, path: dict, start_idx: int) -> List[int]:
        """完成剩余位的译码"""
        remaining_bits = []
        
        for bit_idx in range(start_idx, self.N):
            if bit_idx in self.frozen_bits:
                remaining_bits.append(0)
            else:
                # 简化的硬判决
                remaining_bits.append(0)
        
        return remaining_bits
    
    def simulate_ber(self, snr_db_range: np.ndarray, num_trials: int = 1000) -> np.ndarray:
        """
        仿真比特误码率
        
        Args:
            snr_db_range: SNR范围(dB)
            num_trials: 仿真次数
            
        Returns:
            BER数组
        """
        ber_values = []
        
        for snr_db in snr_db_range:
            errors = 0
            total_bits = 0
            
            for _ in range(num_trials):
                # 生成随机消息
                message = np.random.randint(0, 2, self.K)
                
                # 编码
                codeword = self.encode(message)
                
                # BPSK调制
                tx_signal = 2 * codeword - 1
                
                # 添加噪声
                snr = 10 ** (snr_db / 10)
                sigma = np.sqrt(1 / (2 * snr))
                noise = np.random.normal(0, sigma, self.N)
                rx_signal = tx_signal + noise
                
                # 译码
                decoded = self.decode(rx_signal, snr_db)
                
                # 计算错误
                errors += np.sum(message != decoded)
                total_bits += self.K
            
            ber = errors / total_bits
            ber_values.append(ber)
            print(f"SNR={snr_db}dB, BER={ber:.4e}")
        
        return np.array(ber_values)
